#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
多目标优化模块
实现Pareto前沿计算和拥挤度排序
"""
import time

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import List, Tuple
from matplotlib.pylab import mpl
from src.utils import get_experiment_dir
import os

mpl.rcParams['font.sans-serif'] = ['SimHei']  # 添加这条可以让图形显示中文
mpl.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

class MultiObjectiveOptimizer:
    """多目标优化类，实现Pareto前沿计算和拥挤度排序"""
    
    def divide_fronts(self, answer: List) -> List[List]:
        """
        非支配排序：将解集划分为多个Pareto前沿
        
        Args:
            answer: 解集，每个元素包含多个目标值
        
        Returns:
            front: 分层后的解集，每层包含该层的解索引
        """
        S = [[] for _ in range(len(answer))]
        front = [[]]
        n = [0 for _ in range(len(answer))]  # 记录下比该点厉害的点的数量
        
        # 构建支配关系
        for p in range(len(answer)):
            for q in range(len(answer)):
                # 如果p支配q
                if (np.array(answer[p]) <= np.array(answer[q])).all() and (answer[p] != answer[q]): 
                    if q not in S[p]:
                        S[p].append(q)  # 如果q不属于sp将其添加到sp中
                # 如果q支配p
                elif (np.array(answer[p]) >= np.array(answer[q])).all() and (answer[p] != answer[q]):
                    n[p] = n[p] + 1  # 记录下比该点厉害的点的数量
            
            # 前沿初始化
            if n[p] == 0:
                if p not in front[0]:
                    front[0].append(p)
        
        # 分层
        i = 0
        while front[i]: 
            Q = []
            for p in front[i]:
                for q in S[p]:
                    n[q] = n[q] - 1  
                    if n[q] == 0:   # 如果n[q]==0，说明没有点比q点厉害
                        if q not in Q:
                            Q.append(q)
            i = i + 1
            front.append(Q)
            
        del front[len(front) - 1]  # 最后一个是空数组删掉
        return front

    def calculate_crowding_distance(self, answer: List) -> Tuple[List, List]:
        """
        计算拥挤度并排序
        
        Args:
            answer: 解集，每个元素包含多个目标值
            
        Returns:
            front: 分层后的Pareto前沿
            crowder: 按拥挤度排序后的解索引
        """
        crowder = []
        front = self.divide_fronts(answer)
        answerdf = pd.DataFrame(answer)
        # answerdf.columns = ['x', 'y', 'z']
        answerdf.columns = ['x', 'y']

        
        for i in range(len(front)):
            if len(front[i]) > 1:
                tmp = answerdf.loc[front[i]]
                x = tmp.sort_values(by=['x'], ascending=True)['x']  # 由小到大排序
                y = tmp.sort_values(by=['y'], ascending=True)['y']
                # z = tmp.sort_values(by=['z'], ascending=True)['z']
                tmpx = x.tolist()
                tmpy = y.tolist()
                # tmpz = z.tolist()
                
                # 计算三个维度的拥挤度
                for j in range(1, (tmp.shape[0] - 1)):
                    if tmpx[-1] == tmpx[0]:
                        x.iloc[j] = 0    
                    else:
                        x.iloc[j] = (tmpx[j+1] - tmpx[j-1]) / (tmpx[-1] - tmpx[0])
                    
                    if tmpy[-1] == tmpy[0]:
                        y.iloc[j] = 0
                    else:
                        y.iloc[j] = (tmpy[j+1] - tmpy[j-1]) / (tmpy[-1] - tmpy[0])
                    
                    # if tmpz[-1] == tmpz[0]:
                    #     z.iloc[j] = 0
                    # else:
                    #     z.iloc[j] = (tmpz[j+1] - tmpz[j-1]) / (tmpz[-1] - tmpz[0])
                
                # 端点设置最大拥挤度
                # x.iloc[0], y.iloc[0], z.iloc[0] = 100000, 100000, 100000  # 最左边的点
                # x.iloc[-1], y.iloc[-1], z.iloc[-1] = 100001, 100001, 100001  # 最右边的点

                x.iloc[0], y.iloc[0] = 100000, 100000  # 最左边的点
                x.iloc[-1], y.iloc[-1] = 100001, 100001  # 最右边的点
                
                # 拥挤度由大到小排序
                # f = (x + y + z).sort_values(ascending=False)
                f = (x + y).sort_values(ascending=False)

                crowder = crowder + list(f.index)
            else:
                crowder.append(front[i][0])
                
        return front, crowder

    def calculate_crowding_distance3(self, answer: List) -> Tuple[List, List]:
        """
        计算拥挤度并排序

        Args:
            answer: 解集，每个元素包含多个目标值

        Returns:
            front: 分层后的Pareto前沿
            crowder: 按拥挤度排序后的解索引
        """
        crowder = []
        front = self.divide_fronts(answer)
        answerdf = pd.DataFrame(answer)
        answerdf.columns = ['x', 'y', 'z']
        # answerdf.columns = ['x', 'y']

        for i in range(len(front)):
            if len(front[i]) > 1:
                tmp = answerdf.loc[front[i]]
                x = tmp.sort_values(by=['x'], ascending=True)['x']  # 由小到大排序
                y = tmp.sort_values(by=['y'], ascending=True)['y']
                z = tmp.sort_values(by=['z'], ascending=True)['z']
                tmpx = x.tolist()
                tmpy = y.tolist()
                tmpz = z.tolist()

                # 计算三个维度的拥挤度
                for j in range(1, (tmp.shape[0] - 1)):
                    if tmpx[-1] == tmpx[0]:
                        x.iloc[j] = 0
                    else:
                        x.iloc[j] = (tmpx[j + 1] - tmpx[j - 1]) / (tmpx[-1] - tmpx[0])

                    if tmpy[-1] == tmpy[0]:
                        y.iloc[j] = 0
                    else:
                        y.iloc[j] = (tmpy[j + 1] - tmpy[j - 1]) / (tmpy[-1] - tmpy[0])

                    if tmpz[-1] == tmpz[0]:
                        z.iloc[j] = 0
                    else:
                        z.iloc[j] = (tmpz[j+1] - tmpz[j-1]) / (tmpz[-1] - tmpz[0])

                # 端点设置最大拥挤度
                x.iloc[0], y.iloc[0], z.iloc[0] = 100000, 100000, 100000  # 最左边的点
                x.iloc[-1], y.iloc[-1], z.iloc[-1] = 100001, 100001, 100001  # 最右边的点

                # x.iloc[0], y.iloc[0] = 100000, 100000  # 最左边的点
                # x.iloc[-1], y.iloc[-1] = 100001, 100001  # 最右边的点

                # 拥挤度由大到小排序
                f = (x + y + z).sort_values(ascending=False)
                # f = (x + y).sort_values(ascending=False)

                crowder = crowder + list(f.index)
            else:
                crowder.append(front[i][0])

        return front, crowder

    def draw_objective_changes1(self, fit_every: List) -> None:
        """
        绘制目标函数随迭代次数的变化图（普通调度）
        
        Args:
            fit_every: 每次迭代的目标函数值记录
        """
        plt.figure(figsize=(20, 10))
        font1 = {'family': 'SimHei', 'weight': 'normal', 'size': 22}
        
        legend = [
            ['最小值', '平均值', '最大值'],
            ['最小值', '平均值', '最大值']
        ]
        
        title = ['最大完工时间变化图', '机器负荷极差变化图']
        
        # 只绘制两个子图
        for i in range(2):
            plt.subplot(1, 2, i + 1)  # 改为1行2列布局
            x = [fit_every[i][j][0] for j in range(len(fit_every[i]))]
            y = [fit_every[i][j][1] for j in range(len(fit_every[i]))]
            z = [fit_every[i][j][2] for j in range(len(fit_every[i]))]

            plt.plot(fit_every[3], x, c='black', linestyle='-')
            plt.plot(fit_every[3], y, c='black', linestyle='--')
            plt.plot(fit_every[3], z, c='black', linestyle='-.')

            plt.xlabel('迭代次数', font1)
            plt.title(title[i], font1)
            plt.legend(legend[i], fontsize=18)
            plt.tick_params(labelsize=18)
            
            # 添加网格线以提高可读性
            plt.grid(True, linestyle='--', alpha=0.7)

        plt.tight_layout()  # 优化布局
        
        # 保存到实验目录
        exp_dir = get_experiment_dir()
        save_path = os.path.join(exp_dir, '普通调度多目标优化迭代过程.svg')
        plt.savefig(save_path, bbox_inches='tight')
        print(f"迭代过程图已保存至: {save_path}")

    def draw_objective_changes2(self, fit_every: List) -> None:
        """
        绘制目标函数随迭代次数的变化图（重调度）
        
        Args:
            fit_every: 每次迭代的目标函数值记录
        """
        plt.figure(figsize=(25, 10))
        font1 = {'family': 'SimHei', 'weight': 'normal', 'size': 22}
        
        legend = [
            ['最小值', '平均值', '最大值'],
            ['最小值', '平均值', '最大值'],
            ['最小值', '平均值', '最大值']
        ]
        
        title = ['最大完工时间', '负载均衡', '重调度偏移度']
        
        # 绘制3个子图（重调度模式下维持3个目标）
        for i in range(3):
            plt.subplot(1, 3, i+1)
            x = [fit_every[i][j][0] for j in range(len(fit_every[i]))]
            y = [fit_every[i][j][1] for j in range(len(fit_every[i]))]
            z = [fit_every[i][j][2] for j in range(len(fit_every[i]))]
            
            plt.plot(fit_every[3], x, c='black', linestyle='-')
            plt.plot(fit_every[3], y, c='black', linestyle='--')
            plt.plot(fit_every[3], z, c='black', linestyle='-.')
            
            plt.xlabel('迭代次数', font1)
            plt.title(title[i], font1)
            plt.legend(legend[i], fontsize=18)
            plt.tick_params(labelsize=18)
            
            # 添加网格线
            plt.grid(True, linestyle='--', alpha=0.7)

        plt.tight_layout()
        
        # 保存到实验目录
        exp_dir = get_experiment_dir()
        save_path = os.path.join(exp_dir, '重调度多目标优化迭代过程.svg')
        plt.savefig(save_path, bbox_inches='tight')
        print(f"迭代过程图已保存至: {save_path}") 